from PIL import Image
import sys
import os

def mask_alpha_white(input_path, out_path=None):
    """
    将RGBA图像的Alpha通道设置为白色（完全不透明）
    
    参数:
        input_path: 输入图像路径
        out_path: 输出图像路径（如果为None，则在原文件名后添加_white_alpha）
    
    返回:
        bool: 处理是否成功
    """
    try:
        img = Image.open(input_path).convert("RGBA")
        # 创建白色alpha通道
        white_alpha = Image.new("L", img.size, 255)
        img.putalpha(white_alpha)

        if out_path is None:
            name, ext = os.path.splitext(input_path)
            out_path = f"{name}_white_alpha{ext}"

        img.save(out_path)
        print(f"✓ 处理完成: {input_path} -> {out_path}")
        return True

    except Exception as e:
        print(f"处理失败: {e}")
        print(f"错误信息: {str(e)}")
        return False


def batch_process(folder_path):
    """
    批量处理文件夹中的所有图像文件
    
    参数:
        folder_path: 包含图像文件的文件夹路径
    """
    supported_formats = ('.png', '.tiff', '.tif', '.webp')
    processed_count = 0
    
    for filename in os.listdir(folder_path):
        if filename.lower().endswith(supported_formats):
            input_path = os.path.join(folder_path, filename)
            if mask_alpha_white(input_path):
                processed_count += 1
    
    print(f"\n批量处理完成，成功处理 {processed_count} 个文件")


def main():
    """
    主函数入口，处理命令行参数
    """
    print("=== RGB LUT 图像处理工具 ===\n")
    
    # 检查命令行参数
    if len(sys.argv) == 1:
        # 无参数时，处理当前目录下的所有支持的图像文件
        print("未指定输入，将处理当前目录下的所有图像文件...")
        current_dir = os.getcwd()
        batch_process(current_dir)
    elif len(sys.argv) == 2:
        # 单个参数，可以是文件或目录
        input_path = sys.argv[1]
        if os.path.isfile(input_path):
            # 处理单个文件
            print(f"处理单个文件: {input_path}")
            mask_alpha_white(input_path)
        elif os.path.isdir(input_path):
            # 处理目录
            print(f"批量处理目录: {input_path}")
            batch_process(input_path)
        else:
            print(f"错误: 路径 '{input_path}' 不存在")
            sys.exit(1)
    else:
        print("用法:")
        print("  python extractRGBLut.py          # 处理当前目录下的所有图像文件")
        print("  python extractRGBLut.py <文件>   # 处理单个图像文件")
        print("  python extractRGBLut.py <目录>   # 批量处理目录中的图像文件")
        sys.exit(1)

if __name__ == "__main__":
    main()

